# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/26 9:25 
# @Author  : shutao
# @FileName: stgcn_framework.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from .. import BaseNet


class Graph:
    def __init__(self,
                 layout='openpose',
                 strategy='uniform',
                 max_hop=1,
                 dilation=1):
        self.max_hop = max_hop
        self.dilation = dilation

    def get_edge(self, layout):
        # edge is a list of [child, parent] paris

        if layout:
            self.num_node = 25
            self_link = [(i, i) for i in range(self.num_node)]
            neighbor_link = [(1, 8), (0, 1), (15, 0), (17, 15), (16, 0),
                             (18, 16), (5, 1), (6, 5), (7, 6), (2, 1), (3, 2),
                             (4, 3), (9, 8), (10, 9), (11, 10), (24, 11),
                             (22, 11), (23, 22), (12, 8), (13, 12), (14, 13),
                             (21, 14), (19, 14), (20, 19)]


class STGCN(BaseNet):
    def __init__(self,
                 in_channels=2,
                 edge_importance_weighting=True,
                 data_bn=True,
                 layout='fsd10',
                 strategy='spatial',
                 **kwargs):
        # todo: 修改到此处
        super(STGCN, self).__init__()
        self.data_bn = data_bn
        # load graph
        # self.graph =
        pass

    def forward(self, inputs):
        N, C, T, V, M = inputs.shape
        x = inputs.transpose((0, 4, 3, 1, 2))
        x = x.reshape((N * M, V * C, T))
        if self.data_bn:
            x.stop_gradient = False
        x = self.data_bn(x)

        return inputs

